package com.learning.springboot.trace.zipkin.config;

import brave.CurrentSpanCustomizer;
import brave.SpanCustomizer;
import brave.Tracer;
import brave.Tracing;
import brave.baggage.BaggageField;
import brave.baggage.BaggagePropagation;
import brave.baggage.BaggagePropagationConfig;
import brave.baggage.CorrelationScopeConfig;
import brave.context.slf4j.MDCScopeDecorator;
import brave.handler.SpanHandler;
import brave.http.HttpTracing;
import brave.okhttp3.TracingInterceptor;
import brave.propagation.B3Propagation;
import brave.propagation.CurrentTraceContext;
import brave.propagation.Propagation;
import brave.propagation.ThreadLocalCurrentTraceContext;
import brave.servlet.TracingFilter;
import brave.spring.webmvc.SpanCustomizingAsyncHandlerInterceptor;
import okhttp3.Call;
import okhttp3.OkHttpClient;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Import;
import org.springframework.http.client.OkHttp3ClientHttpRequestFactory;
import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurerAdapter;
import zipkin2.reporter.Reporter;
import zipkin2.reporter.Sender;
import zipkin2.reporter.brave.AsyncZipkinSpanHandler;
import zipkin2.reporter.okhttp3.OkHttpSender;

import javax.servlet.Filter;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * ClassName: SpringbootTraceZipkinApplication
 * Description:
 * Date: 2016/12/22 21:02
 *
 * @author SAM SHO
 * @version V1.0
 */
@Configuration
@Import(SpanCustomizingAsyncHandlerInterceptor.class)
public class TracingAutoConfiguration {

    static final BaggageField USER_NAME = BaggageField.create("userName");

    /**
     * 允许日志模式使用 {@code %{traceId}} {@code %{spanId}} 和 {@code %{userName}}
     * <p>
     * Allows log patterns to use {@code %{traceId}} {@code %{spanId}} and {@code %{userName}}
     */
    @Bean
    CurrentTraceContext.ScopeDecorator correlationScopeDecorator() {
        return MDCScopeDecorator.newBuilder()
                .add(CorrelationScopeConfig.SingleCorrelationField.create(USER_NAME)).build();
    }

    /**
     * Trace 上下文  TraceContext
     * 在线程之间传播跟踪上下文
     * Propagates trace context between threads.
     */
    @Bean
    CurrentTraceContext currentTraceContext(CurrentTraceContext.ScopeDecorator correlationScopeDecorator) {
        return ThreadLocalCurrentTraceContext.newBuilder()
                .addScopeDecorator(correlationScopeDecorator)
                .build();
    }

    /**
     * 配置传播 {@link #USER_NAME}，使用远程头 "USER_NAME"
     * <p>
     * Configures propagation for {@link #USER_NAME}, using the remote header "user_name"
     */
    @Bean
    Propagation.Factory propagationFactory() {
        return BaggagePropagation.newFactoryBuilder(B3Propagation.FACTORY)
                .add(BaggagePropagationConfig.SingleBaggageField.newBuilder(USER_NAME).addKeyName("user_name").build())
                .build();
    }

    /**
     * 配置如何发送跨越到Zipkin，支持多种
     * <p>
     * Configuration for how to send spans to Zipkin
     */
    @Bean
    Sender sender(
            @Value("${zipkin.baseUrl:http://127.0.0.1:9411}/api/v2/spans") String zipkinEndpoint) {
        return OkHttpSender.create(zipkinEndpoint);
    }

    /**
     * 配置如何为Zipkin的消息缓冲跨越（代替 reporter）
     * 默认是 LogSpanHandler
     * <p>
     * Configuration for how to buffer spans into messages for Zipkin
     */
    @Bean
    SpanHandler spanHandler(Sender sender) {
        // 日志打印
        // return AsyncZipkinSpanHandler.create(new LoggingReporter());


        // 异步发送
        return AsyncZipkinSpanHandler.create(sender);
    }

    /**
     * 被 LogSpanHandler 替代
     */
    static final class LoggingReporter implements Reporter<zipkin2.Span> {
        final Logger logger = Logger.getLogger(Tracer.class.getName());

        @Override
        public void report(zipkin2.Span span) {
            if (span == null) {
                throw new NullPointerException("span == null");
            }
            if (!logger.isLoggable(Level.INFO)) {
                return;
            }
            logger.info(span.toString());
        }

        @Override
        public String toString() {
            return "LoggingReporter{name=" + logger.getName() + "}";
        }
    }


    /**
     * 控制跟踪的各个方面，如在UI中显示的服务名称
     * <p>
     * Controls aspects of tracing such as the service name that shows up in the UI
     */
    @Bean
    Tracing tracing(
            @Value("${brave.localServiceName:${spring.application.name}}") String serviceName,
            @Value("${brave.supportsJoin:true}") boolean supportsJoin,
            @Value("${brave.traceId128Bit:false}") boolean traceId128Bit,
            Propagation.Factory propagationFactory,
            CurrentTraceContext currentTraceContext
            , SpanHandler spanHandler
    ) {
        return Tracing.newBuilder()
                .localServiceName(serviceName)
                // .localIp()
                // .localPort()
                .supportsJoin(supportsJoin)
                .traceId128Bit(traceId128Bit)
                // 传播工厂
                .propagationFactory(propagationFactory)
                //
                .currentTraceContext(currentTraceContext)
                // 默认是 LogSpanHandler
                // .addSpanHandler(spanHandler)
                .build();
    }

    /**
     * 允许某人在跟踪过程中向span添加标记。
     * <p>
     * Allows someone to add tags to a span if a trace is in progress.
     */
    @Bean
    SpanCustomizer spanCustomizer(Tracing tracing) {
        return CurrentSpanCustomizer.create(tracing);
    }

    /**
     * 决定如何命名和标记跨度。默认情况下，它们的命名与http方法相同。
     * <p>
     * Decides how to name and tag spans. By default they are named the same as the http method.
     */
    @Bean
    HttpTracing httpTracing(Tracing tracing) {
        return HttpTracing.create(tracing);
    }

    /**
     * 为HTTP请求创建服务器跨度
     * <p>
     * Creates server spans for HTTP requests
     */
    @Bean
    Filter tracingFilter(HttpTracing httpTracing) {
        return TracingFilter.create(httpTracing);
    }

    /**
     * 将MVC控制器标签添加到服务器范围
     * <p>
     * Adds MVC Controller tags to server spans
     */
    @Bean
    WebMvcConfigurer tracingWebMvcConfigurer(
            final SpanCustomizingAsyncHandlerInterceptor webMvcTracingCustomizer) {
        return new WebMvcConfigurerAdapter() {
            /**
             * 将应用程序定义的web控制器详细信息添加到HTTP服务器范围
             * Adds application-defined web controller details to HTTP server spans
             */
            @Override
            public void addInterceptors(InterceptorRegistry registry) {
                registry.addInterceptor(webMvcTracingCustomizer);
            }
        };
    }

    /**
     * 为HTTP请求创建客户端跨度。
     * <p>
     * Creates client spans for HTTP requests.
     *
     * <p>{@link OkHttpClient} because {@link OkHttp3ClientHttpRequestFactory} doesn't take a
     * {@link Call.Factory}
     */
    @Bean
    OkHttpClient tracedOkHttpClient(HttpTracing httpTracing) {
        return new OkHttpClient.Builder()
                .addNetworkInterceptor(TracingInterceptor.create(httpTracing))
                .build();
    }

}